########################################
# Survey 4 Functions
########################################

# packages
if (!require("pacman")) install.packages("pacman")
pacman::p_load(ggplot2,
               bda,
               foreign,
               xtable,
               reshape,
               psych,
               haven,
               systemfit,
               lfe,
               pwr,
               reshape2,
               dplyr,
               magrittr,
               testthat,
               stargazer,
               NPC)

########################################
# Data Processing Functions

# adding up columns with NAs
psum <- function(..., na.rm=T) { 
  x <- list(...)
  rowSums(matrix(unlist(x), ncol=length(x)), na.rm=T)
} 

# function to relabel
relab <- function(old.var, old.labs, new.labs, text = FALSE){
  map <- setNames(object = old.labs, nm = new.labs)
  if (text) {
    names(map[unlist(eval(old.var))])
  } else {
    as.numeric(names(map[unlist(eval(old.var))])) 
  }
}

# function to standardize
vscale <- function(var, vig){
  sv <- c()
  for (i in 1:length(var)){
    v <- vig[i]
    sv[i] <- (var[i] - mean(var[vig==v], na.rm = T))/
      sd(var[vig==v], na.rm=T)
  }
  return(sv)
}

# function to organize multiple versions of the same question
multiq <- function(v1, v2, v3, v4, v5, v6){
  res <- rep(NA, length(v1))
  res[!is.na(eval(v1))] <- eval(v1)[!is.na(eval(v1))]
  res[!is.na(eval(v2))] <- eval(v2)[!is.na(eval(v2))]
  res[!is.na(eval(v3))] <- eval(v3)[!is.na(eval(v3))]
  return(res)
}

# function to sort out the labels for regions
region.lab <- function(region1, region2){
  t <- ifelse(is.na(region1), region2, region1)
  s <- c()
  s[t == 1] <- FALSE
  s[t == 0] <- TRUE
  s[t == -99] <- FALSE
  return(s)
}
#####################################
# Analysis Functions

# robust SE function
robustse <- function(model){ 
  # OLS with robust SE
  # written by Baobao Zhang
  require("sandwich")
  require("lmtest")
  model$newse<-vcovHC(model,type = "HC3")
  coef <- coeftest(model,model$newse)
  return(as.matrix(coef))
}

# function to perform OLS regression and extract coef.
myreg <- function(Y, Z=d$Z, V=d$V, row=2, print.out=FALSE){
  Y <- as.numeric(eval(Y))
  Z <- as.numeric(eval(Z))
  V <- as.numeric(eval(V))
  res <- matrix(NA, max(V, na.rm = TRUE), 2)
  for (i in 1:max(V, na.rm = TRUE)){
    regout <- robustse(lm(Y[V==i]~Z[V==i]))
    res[i,] <- regout[row,1:2]
    if (print.out) {print(regout)}
  }
  return(res)
}

# function to perform 2SLS regression
manualiv <- function (instrument,treatment,outcome){
  # written by Baobao Zhang
  first <- lm(treatment ~ instrument) # first stage
  treatment.hat <- first$fitted
  second <- lm(outcome ~ treatment.hat)
  X.adj <- cbind(rep(1, length(treatment)), treatment)
  second$residuals <- outcome - X.adj %*% coef(second)
  return(robustse(second))
}

# function to perform 2SLS regression and extract the coefficients
myreg.iv <- function(Y, D=d$D, Z=d$Z, V=d$V, row=2, print.out=FALSE){
  Y <- as.numeric(eval(Y))
  Z <- as.numeric(eval(Z))
  V <- as.numeric(eval(V))
  D <- as.numeric(eval(D))
  res <- matrix(NA, max(V, na.rm = TRUE), 2)
  for (i in 1:max(V, na.rm = TRUE)){
    regout <- manualiv(instrument = Z[V==i], 
                       treatment = D[V==i], outcome = Y[V==i])
    res[i,] <- regout[row,1:2]
    if (print.out) {print(regout)}
  }
  return(res)
}

#########################################
# Graphics Functions

# theme for graphics
theme_bb <- function (base_size = 12, base_family = "",where="bottom") 
{
  theme_grey(base_size = base_size, base_family = base_family) %+replace% 
    theme(axis.text = element_text(size = rel(0.8)), 
          legend.position=where, 
          axis.ticks = element_line(colour = "black"), 
          legend.key = element_rect(colour = "grey80"), 
          panel.background = element_rect(fill = "white", colour = NA), 
          panel.border = element_rect(fill = NA,colour = "grey50"), 
          panel.grid.major = element_line(colour = "grey90", size = 0.2), 
          panel.grid.minor = element_line(colour = "grey98", size = 0.5), 
          strip.background = element_rect(fill = "grey80",  colour = "grey50"), strip.background = element_rect(fill = "grey80", 
                                                                                                                colour = "grey50"))
}

# Plot density function
get_density <- function(input_data, input_subset, bandwidth) {
  input_subset <- as.character(unlist(input_subset))
  input_data <- input_data[as.character(input_data$Z) == input_subset[1] &
                             as.character(input_data$V) == input_subset[2] &
                             as.character(input_data$variable) == input_subset[3],]
  xhist <- binning(input_data["value"], 
                   breaks = ifelse(round(min(input_data["value"], 
                                             na.rm = TRUE))-1 > 0,
                                   round(min(input_data["value"], 
                                             na.rm = TRUE))-1, 0), 
                   bw = bandwidth)
  den <- bda::bde(x=xhist, type="smkde", 
                  from = round(min(input_data["value"], na.rm = TRUE)) - 1,
                  to = round(max(input_data["value"], na.rm = TRUE)) +1)
  temp <- data.frame(Z = input_subset[1], 
                     V = input_subset[2],
                     variable = input_subset[3],
                     x = den$x, y = den$y)  
  return(temp)
}

myden <- function(plot_data, combo, short_name, xlab, bandwidth,
                  glab=c("Non-democracy", "Democracy"),
                  color1="orange1", color1b="orange3", bw = 1,
                  color2="steelblue4", ylab="Density",
                  alpha=0.75,  
                  x.limits = NULL, x.breaks = NULL, x.labels = NULL,
                  print_colormodel = "gray"){
  d_res <- get_density(input_data = plot_data, input_subset = combo[1,], 
                       bandwidth = bandwidth)
  for (i in 2:nrow(combo)) {
    d_res <- rbind(d_res, get_density(input_data = plot_data, 
                                      input_subset = combo[i,],
                                      bandwidth = bandwidth))
  }
  # add in means
  summary_data <- plot_data %>% group_by(Z, V, variable) %>% 
    dplyr::summarise(mean_value = mean(value, na.rm = TRUE)) %>% as.data.frame()
  d_res <- merge(x = d_res, y = summary_data, all.x = TRUE)
  # fix factors
  d_res$Z <- factor(d_res$Z, levels = levels(plot_data$Z))
  d_res$variable <- factor(d_res$variable, levels = levels(plot_data$variable))
  d_res$V <- factor(d_res$V, levels = levels(plot_data$V))
  # generate plot
  if (is.null(x.breaks) & is.null(x.labels) & is.null(x.limits)) {
    my_plot <- ggplot(d_res, 
           aes(x = x, y = y, fill = Z)) + 
      geom_area(position = "identity", alpha = 0.75, color = "black") + 
      facet_grid(V ~ .) + theme_bb() +
      geom_vline(aes(xintercept = mean_value, color = Z), 
                 linetype="longdash", size=1) +
      ylab("Density") + xlab(xlab) +
      scale_color_manual(name="Treatment Assignment Regime Type", 
                         values=c(color1, color2)) + 
      scale_fill_manual(name="Treatment Assignment Regime Type", 
                        values=c(color1, color2))
  } else {
    my_plot <- ggplot(d_res, 
                      aes(x = x, y = y, fill = Z)) + 
      geom_area(position = "identity", alpha = 0.75, color = "black") + 
      facet_grid(V ~ .) + theme_bb() +
      geom_vline(aes(xintercept = mean_value, color = Z), 
                 linetype="longdash", size=1) +
      ylab("Density") + xlab(xlab) +
      scale_color_manual(name="Treatment Assignment Regime Type", 
                         values=c(color1, color2)) + 
      scale_fill_manual(name="Treatment Assignment Regime Type", 
                        values=c(color1, color2)) +
      scale_x_continuous(limits=x.limits, breaks=x.breaks,
                         labels=x.labels)
  }
  print(my_plot)
  ggsave(paste0(images_directory, "/", short_name, ".pdf"), 
         width=6, height = 7, dpi = 300, colormodel = print_colormodel) 
}

# Plot histograms
myhist <- function(X, Y, glab=c("Non-democracy", "Democracy"), 
                   color1="orange1", color2="steelblue4",
                   outline="black", mytheme=theme_bb(),
                   title, xlab, ylab){
  data <- na.omit(data.frame(eval(X), eval(Y)))
  data$x <- factor(data[,1], labels=glab)
  data$y <- factor(data[,2], labels=c("No", "Yes"))
  ggplot(data=data, aes(x=y, fill=x))+
    geom_histogram(position="dodge", color=outline)+
    xlab(paste0(xlab, "\n N=", length(Y[V==v & !is.na(Y)]))) +
    ylab(ylab)+ mytheme+ ggtitle(title) +
    scale_fill_manual(values=c(color1, color2),
                      name="Treatment Assignment Regime Type")
}

# ggplot2 themes


theme_bb <- function (base_size = 12, base_family = "",where="bottom") 
{
  theme_grey(base_size = base_size, base_family = base_family) %+replace% 
    theme(axis.text = element_text(size = rel(0.8)), 
          legend.position=where, 
          axis.ticks = element_line(colour = "black"), 
          legend.key = element_rect(colour = "grey80"), 
          panel.background = element_rect(fill = "white", colour = NA), 
          panel.border = element_rect(fill = NA,colour = "grey50"), 
          panel.grid.major = element_line(colour = "grey90", size = 0.2), 
          panel.grid.minor = element_line(colour = "grey98", size = 0.5), 
          strip.background = element_rect(fill = "grey80", colour = "grey50"))
}


theme_bare <- function (base_size = 12, base_family = "") 
{
  theme_grey(base_size = base_size, base_family = base_family) %+replace% 
    theme(axis.text = element_text(size = rel(0.8)), 
          legend.position="none", 
          axis.ticks = element_line(colour = "black"), 
          legend.key = element_rect(colour = "grey80"), 
          panel.background = element_rect(fill = "white", colour = NA), 
          panel.border = element_rect(fill = NA,colour = "grey50"), 
          panel.grid.major = element_line(colour = "grey90", size = 0.2), 
          panel.grid.minor = element_line(colour = "grey98", size = 0.5), 
          strip.background = element_rect(fill = "grey80", colour = "grey50"))
}

theme_bb3 <- function (base_size = 12, base_family = "") 
{
  theme_grey(base_size = base_size, base_family = base_family) %+replace% 
    theme(axis.text = element_text(size = rel(0.8)), 
          legend.position="right", 
          axis.ticks = element_line(colour = "black"), 
          legend.key = element_rect(colour = "grey80"), 
          panel.background = element_rect(fill = "white", colour = NA), 
          panel.border = element_rect(fill = NA,colour = "grey50"), 
          panel.grid.major = element_line(colour = "grey90", size = 0.2), 
          panel.grid.minor = element_line(colour = "grey98", size = 0.5), 
          strip.background = element_rect(fill = "grey80", colour = "grey50"))
}


theme_bb4 <- function (base_size = 12, base_family = "", height, width) 
{
  theme_grey(base_size = base_size, base_family = base_family) %+replace% 
    theme(axis.text = element_text(size = rel(0.8)), 
          legend.position="right", 
          axis.ticks = element_line(colour = "black"), 
          legend.key = element_rect(colour = "grey80"), 
          panel.background = element_rect(fill = "white", colour = NA), 
          panel.border = element_rect(fill = NA,colour = "grey50"), 
          panel.grid.major = element_line(colour = "grey90", size = 0.2), 
          panel.grid.minor = element_line(colour = "grey98", size = 0.5), 
          strip.background = element_rect(fill = "grey80",  colour = "grey50"), 
          legend.key.height = unit(height, "cm"),
          legend.key.width = unit(width, "cm"))
}

theme_bbtop <- function (base_size = 12, base_family = "", height, width,
                         legend_position = "top",
                         stack_direction = "horizontal") 
{
  theme_grey(base_size = base_size, base_family = base_family) %+replace% 
    theme(axis.text = element_text(size = rel(0.8)), 
          legend.position=legend_position, 
          legend.direction=stack_direction, 
          axis.ticks = element_line(colour = "black"), 
          legend.key = element_rect(colour = "grey80"), 
          panel.background = element_rect(fill = "white", colour = NA), 
          panel.border = element_rect(fill = NA,colour = "grey50"), 
          panel.grid.major = element_line(colour = "grey90", size = 0.2), 
          panel.grid.minor = element_line(colour = "grey98", size = 0.5), 
          strip.background = element_rect(fill = "grey80", colour = "grey50"))
}

theme_nolegend <- function (base_size = 12, base_family = "", height, width) 
{
  theme_grey(base_size = base_size, base_family = base_family) %+replace% 
    theme(axis.text = element_text(size = rel(0.8)), 
          legend.position="none", 
          axis.ticks = element_line(colour = "black"), 
          legend.key = element_rect(colour = "grey80"), 
          panel.background = element_rect(fill = "white", colour = NA), 
          panel.border = element_rect(fill = NA,colour = "grey50"), 
          panel.grid.major = element_line(colour = "grey90", size = 0.2), 
          panel.grid.minor = element_line(colour = "grey98", size = 0.5), 
          strip.background = element_rect(fill = "grey80", colour = "grey50"))
}

# ggplot 
coef_plot_function <- function(start_ggplot, file_name, x_title_name,
                               scale_type = "fixed", 
                               scale_name = "Vignette Type",
                               stack_direction = "horizontal",
                               legend_position = "top",
                               color_values = 
                                 c("firebrick3","forestgreen","dodgerblue3"),
                               main_title_name, factet_yes = TRUE, 
                               shape_values = c(21,22,23),
                               out_height, out_width, alpha_values,
                               dot_size = 4, print_out = TRUE, 
                               print_colormodel = "grey") {
  f <- start_ggplot + geom_vline(xintercept=0, linetype="longdash")+
    geom_errorbarh(aes(xmax =  coef + qnorm(1 - 0.005)*se, 
                       xmin = coef + qnorm(0.005)*se),
                   size=0.6, height=0) +
    geom_errorbarh(aes(xmax =  coef + qnorm(1 - 0.025)*se, 
                       xmin = coef + qnorm(0.025)*se),
                   size=1.0, height=0)+
    geom_point(stat="identity", fill="white", size = dot_size)+
    scale_color_manual(name = scale_name,
                       values=color_values)+
    scale_shape_manual(name=scale_name, values = shape_values) +
    scale_alpha_manual(name = scale_name, values = alpha_values) + 
    theme_bbtop(stack_direction = stack_direction, 
                legend_position = legend_position) + 
    ggtitle(main_title_name) + 
    xlab(x_title_name)+
    ylab("") + scale_y_discrete(breaks=NULL) 
  if (factet_yes){
    f <- f + facet_wrap( ~ placebo, ncol=1, scales = scale_type) 
  }
  if (print_out) {
    print(f)
    if (!is.null(images_directory)) {
      ggsave(paste0(images_directory, "/", file_name, ".pdf"), 
             height=out_height, width=out_width, dpi = 300, 
             colormodel = print_colormodel)
    }  
  } else {
    return(f)
  }
}

